open Num
open List


(* Util *)

let (%) f g x = g (f x)
let (%.) f g x y = f (g x y)

let cons x xs = x :: xs
let pair a b = (a,b)
let rec from_to a b = if a <= b then a :: from_to (a+1) b else []

let unique = sort compare % 
  (let rec unique = function
     | [] -> []
     | x :: x' :: xs when x = x' -> unique (x :: xs)
     | x :: xs -> x :: unique xs
   in unique)

let group (xs:('a*'b) list) : ('a * ('b * 'b list)) list =
  let h = Hashtbl.create (length xs) in
  iter (fun (a,b) -> Hashtbl.add h a b) xs;
  let keys = (map fst  % unique) xs in
    map (fun a -> a, 
         match Hashtbl.find_all h a with
           | b :: bs -> (b,bs)
           | [] -> assert false) keys

let rec ix i = function 
  | [] -> []
  | x :: xs -> (i,x) :: ix (i+1) xs

let rec get_ix y = function
  | x::xs -> if x = y then 0 else 1 + get_ix y xs
  | [] -> raise Not_found

let rec create a = function 
  | 0 -> []
  | n -> a :: create a (n-1)

let rec prefix i = function
  | _ when i = 0 -> []
  | x::xs -> x :: prefix (i-1) xs
  | [] -> failwith "prefix"

let fold_left' f x0 = function
  | x::xs -> fold_left f x xs
  | [] -> x0

let ( /%/ ) f g = g % f

let num_of_float_string s =
  let ix_p = String.index s '.' in
  let ix_e = String.index_from s (ix_p+1) 'E' in
  let l = ix_e-ix_p-1 in
    (num_of_string (String.sub s 0 ix_p) +/
    (num_of_string (String.sub s (ix_p+1) l) // (num_of_int 10 **/ num_of_int l))) */
    (num_of_int 10 **/ (num_of_string (String.sub s (ix_e+1) (String.length s - ix_e - 1))))

let round_num' a x = round_num (x//a) */ a

let float_string_of_num = approx_num_fix 20

let error_msg fn (l,exc) =
  let (_,l,c1,c2) = Stdpp.line_of_loc fn l in
  prerr_endline ("\nFile \"" ^ fn ^ "\", line " ^ string_of_int l ^
    ", characters " ^ string_of_int c1 ^ "-" ^ string_of_int c2 ^ ":");
  match exc with
      | Stream.Error msg | Failure msg -> prerr_endline msg
      | exc -> raise exc


let zero = num_of_int 0
let one = num_of_int 1


(* Poly *)

type mono = int list
type poly =  (mono * num) list

let rec chop a = function
  | x :: xs when x=a -> chop a xs
  | xs -> xs

let normalize_mono = rev % chop 0 % rev

let rec mono_mul xs ys = match xs, ys with
  | x :: xs', y :: ys' -> (x+y)::mono_mul xs' ys'
  | [], ys -> ys
  | xs, [] -> xs

let mono_div m1 m2 = mono_mul m1 (map (fun x -> -x) m2)

let poly_var i : poly = [ create 0 i @ [1], one ]

let rec poly_const c : poly = [[], c]

let rec poly_add (xs:poly) : poly -> poly = function
  | [] -> xs 
  | (y,c) :: ys ->
      match partition (fun (x,_) -> x=y) xs with
        | ([],_) -> (y,c) :: poly_add xs ys
        | ([x,c2],xs') -> (y,c +/ c2) :: poly_add xs' ys
        | (_,_) -> failwith "add: second argument not in normal form"

let rec poly_mul (xs:poly) : poly -> poly = function
  | [] -> []
  | (y,c) :: ys ->
      poly_add
      (map (fun (x,c2) -> mono_mul x y, c*/c2) xs)
      (poly_mul xs ys)

let rec poly_neg : poly -> poly = map (fun (m,c) -> (m, minus_num c)) 

let degree : poly -> int = fold_left max 0 /%/ map (fold_left (+) 0 /%/ fst)

let rec poly_pow p = function
  | 0 -> poly_const one
  | n -> poly_mul p (poly_pow p (n-1))

let normalize_poly : poly -> poly =
  group %
  map (fun (m,(c,cs)) -> m, fold_left' (+/) c cs) %
  filter (snd % (=/) zero % not) %
  sort (fun (m1,_) (m2,_) -> compare m1 m2)

let rec map2' f xs ys = match xs, ys with
  | _, [] -> []
  | x::xs, y::ys -> f x y :: map2' f xs ys
  | _ -> failwith "map2'"

let string_of_mono : string list -> mono -> string = 
  (fun xs -> let r = fold_left' (fun a b -> a ^ "*" ^ b) "1" xs in 
     if length xs > 1 then "(" ^ r ^ ")" else r) /%/
  filter ((<>) "")
    %. 
    map2' (fun x i -> match i with
            | 0 -> ""
            | 1 -> x
            | _ -> x ^ "^" ^ string_of_int i) 


let string_of_poly x : poly -> string = 
  normalize_poly %
    map (function
           | (m,c) when c =/ one -> string_of_mono x m
           | (m,c) -> string_of_num c ^ "*" ^ string_of_mono x m
         ) %
  fold_left' (fun a b -> a ^ " + " ^ b) "0" 

let rec monos_up_to d = function
  | 0 -> [ [] ]
  | n -> flatten (map (fun e -> map (cons e) (monos_up_to (d-e) (n-1))) (from_to 0 d))

let monos_up_to d = monos_up_to d % map normalize_mono


(* Matrix *)

let sym x = map (fun ((i,j),c) -> (i,j), if i = j then c else c // num_of_int 2) x

type dense = num list list
type dense_sym = num list list
type low_tri = num list list
type sparse = ((int * int) * num) list
type sparse_sym = ((int * int) * num) list
type diag = num list

let rec transpose = function (* for dense_sym *)
  | [] -> []
  | m -> map hd m :: transpose (tl (map tl m))

let matrix_add : dense -> dense -> dense = map2 (map2 (+/))
let matrix_mul_cons (c : num) : dense -> dense = map (map (( */ ) c))

let dense_of_dense_sym m = map2 (@) m (map tl (transpose m))

let dense_sym_of_sparse_sym k (s:sparse_sym) : dense_sym =
  let d = Array.make_matrix k k zero in
  iter (fun ((i,j),c) -> d.(max i j-1).(min i j-1) <- c) s;
  (Array.to_list /%/ Array.mapi (fun i r -> prefix (i+1) (Array.to_list r))) d

let dense_of_lowtri m = 
  map2 (fun r k -> r @ create zero k) m (rev (from_to 0 (length m-1)))

let dense_of_diag xs = 
  let zs = map (create zero) (from_to 0 (length xs - 1)) in
    map2 (fun x (z1,z2) -> (z1@x::z2)) xs (combine zs (rev zs))

let dot_prod = fold_left (+/) zero %. map2 ( */ )

let (/.%) f g x = let (a,b) = f x in (g a, b)

exception Not_psd

let rec diag_row m r ai : (num*num) list * num = match m, ai with
  | (ld_l_j,d_j)::m', ai0::ai' -> 
      let ld_ji = ai0 -/ dot_prod (map fst ld_l_j) (map snd r) in
        diag_row m' (r @ [ ld_ji, 
                           if d_j =/ zero then
                             (if ld_ji =/ zero then Int 42 (* any *) else raise Not_psd)
                           else ld_ji // d_j 
                         ]) ai'
  | [], [aii] -> r, aii -/ dot_prod (map fst r) (map snd r)
  | _, _ -> failwith "diag_row"

let diag : dense_sym -> (low_tri * diag) option = fun x -> 
  try
    Some (
      (fold_left (fun l r -> l @ [ diag_row l [] r ]) []
       % split /.% 
       map (map snd % (fun r -> r @ [one]))
      ) x)
  with Not_psd -> None


let sparse_diag n (a:sparse) =
  let rec sd i = 
    if i = n then [] else
      (try assoc (i,i) a
       with Not_found -> zero) :: sd (i+1)
  in sd 0


(* Note: This diagonalization algorithm is designed to yield small
   expressions as matrix entries. *)


(* Syntax *)

let gram = Grammar.gcreate (Plexer.gmake())
let term = Grammar.Entry.create gram "term"
let problem = Grammar.Entry.create gram "problem"

let result = Grammar.Entry.create gram "result"
let entry = Grammar.Entry.create gram "entry"
let sci_float = Grammar.Entry.create gram "sci_float"

let monad1 f a x = 
  let (a',x') = a x in
    f a', x'

let monad2 f a b x = 
  let (a',x') = a x in
  let (b',x'') = b x' in
    f a' b', x''

EXTEND
  problem:  [[ t = term; EOI -> t [] ]];

  term:
    [ [ a = term; "+"; b = term -> monad2 poly_add a b
      | a = term; "-"; b = term -> monad2 poly_add a (monad1 poly_neg b)
      | "-"; a = term -> monad1 poly_neg a
      ]
    | [ a = term; "*"; b = term -> monad2 poly_mul a b
      ]
    |  [  f = LIDENT; i = OPT [ "!"; i = INT -> i ]; args = LIST0 NEXT -> fun vars ->
            let (args,vars) = fold_left (fun (args',vars) a -> 
                                           let (a',vars') = a vars in (args' @ [a']), vars') ([],vars) args in
            let v = fold_left (fun a b -> a ^ " (" ^ string_of_poly vars b ^ ")") (f^(match i with Some i -> "!" ^ i | None -> "")) args in
              (try poly_var (get_ix v vars), vars
               with Not_found -> 
                 poly_var (length vars), (vars @ [ v ])
              )
    ]|[ x = LIDENT; i = OPT [ "!"; i = INT -> i ] -> fun vars ->
       let v = x ^ (match i with Some i -> "!" ^ i | None -> "") in
        (try poly_var (get_ix v vars), vars
               with Not_found -> 
                 poly_var (length vars), (vars @ [ v ]))
      | c = FLOAT -> fun vars -> poly_const (num_of_string c), vars
      | c1 = INT; c2 = OPT ["/"; i = INT -> i]-> fun vars -> poly_const (num_of_string (match c2 with None -> c1 | Some c2 -> c1 ^ "/" ^ c2)), vars

      | "("; t = term; ")" -> t
      ]
    | [ a = term; "^"; n = INT -> let n = int_of_string n in
          if n < 0 then Stdpp.raise_with_loc _loc (Failure "negative exponent") else
            monad1 (fun x -> poly_pow x n) a
      ]
];

  result: [[ c = LIST0 sci_float; m = LIST0 entry; EOI -> c, m ]];
  entry: [[ id=INT; block=INT; i=INT; j=INT; c=sci_float -> ()]];
  sci_float: [[ m=OPT "-"; x=FLOAT -> 
                  round_num' (num_of_string "1/1000") (* experimental choice *)
              (match m with Some _ -> minus_num (num_of_float_string x) | _ -> num_of_float_string x) ]];
END


(* S.O.S *)

let submonos n (p:poly) = filter
    (fun m -> exists (fst % mono_div m % for_all ((>=) 0)) p)
    (monos_up_to (degree p / 2) n)

let rec cart_up_tri f = function
  | [] -> []
  | x::xs -> map (f x) (x::xs) :: cart_up_tri f xs

let submono_prods sm =
  (cart_up_tri (fun (i1,m1) (i2,m2) -> mono_mul m1 m2, (i1,i2)) % flatten % group) (ix 1 sm)

let rec ungroup = function
  | [] -> []
  | (x,ys)::xs -> map (pair x) ys @ ungroup xs

let sdp_of_sos sm p : int * sparse_sym * sparse_sym list =
  let smp = submono_prods sm in
    prerr_endline ((length % string_of_int) smp ^ " submonomial products"); flush stderr;
(  length sm,
  (map (fun (m,c) -> fst (assoc m smp),  minus_num c) % sym) p,
    (filter (fun (_,(_,ms)) -> ms <> []) %
     map snd %
       ungroup
     %
     map (fun ((m,m')) ->
            assert (m<>m');
          sym [m, Int (-1); m', Int 1]
         )
    ) smp
)


let format_sdpa (k,f,fs) out =
  let sparse (id,cs) =
    let pre = string_of_int id ^ " 1 " in
    fold_left (fun out ((i,j),c) -> out ^ "\n" ^ pre ^ string_of_int i ^ " " ^ 
                 string_of_int j ^ " " ^ float_string_of_num c) "\n" cs in
  let m = length fs in
    output_string out (string_of_int m  ^ "\n1\n" ^ string_of_int k ^ "\n" ^  (fold_left (^) "" (create "1. " m)));
    iter (output_string out) (map sparse (ix 0 (f::fs)))
    
(* alternatives for diag:
- (fold_left (fun a b -> a ^ " " ^ float_string_of_num b) "" /%/ map (fun d -> d*/d))
  (map (fold_left (+/) zero /%/ sparse_diag k) fs)
- (fold_left (^) "" (create "1. " m))
- "1. " ^  fold_left (^) "" (create "0. " (m-1))
*)

let solve_sdp (k,(f0:sparse_sym),(fs:sparse_sym list)) : (num list (* [ dense_sym]*)) option = 
  let fn_out = "p.dat-s" in
  let fn_in = "p.out" in
  let out = open_out fn_out in
    prerr_endline "Writing SDP problem..."; flush stderr;
    (format_sdpa (k,f0,fs) out); close_out out;
    prerr_endline "Calling csdp..."; flush stderr;
    let r = Sys.command ("csdp " ^ fn_out ^ " " ^ fn_in (*^ " > /dev/null"*)) in
          prerr_endline "Csdp returned.";flush stderr;
  if r = 0 || r = 3 then (* 0 means success, 3 means "partial" success *)
    Some (
      let inp = open_in fn_in in
      let (c,xy) = try Grammar.Entry.parse result (Stream.of_channel inp) 
      with Stdpp.Exc_located (loc,exc) -> (error_msg fn_in (loc,exc); exit 1) in
        (*  let xy = assoc 1 (group' xy) in*)
        close_in inp; c(*, xy*))
  else
    None


let zld z l d : (num * poly) list = 
  filter (fst % (=/) zero % not)
    (combine d (map (map2' pair (rev z)) ((map rev /%/ transpose) l)))

let zld_to_string x z l d = 
  fold_left' (fun a b -> a ^ " +\n" ^ b) "" 
  (map (fun (c,p) -> string_of_num c ^ "*(" ^ string_of_poly x p ^ ")^2")  (zld z l d))

let rec disturb (eps : num) : num list -> num list list = function
  | [] -> [[]]
  | x :: xs when x =/ Int 0 -> prerr_string "."; flush stderr;
      map (cons (Int 0)) (disturb eps xs) (* assume that 0 is always exact *)
  | x :: xs -> 
      let xs' = disturb eps xs in
        map (cons x) xs' @ 
        map (cons (x*/(one+/eps))) xs' @ 
        map (cons (x*/(one-/eps))) xs'

let lin_comb l f0 fs = 
  fold_left matrix_add
    (matrix_mul_cons (minus_num one) f0)
    (map2 matrix_mul_cons l fs)

let ma_of_row r =
  "{" ^ fold_left' (fun a b -> a ^ ", " ^ b) "" (map string_of_num r) ^ "}"

let ma_of_dense (d:dense) =
  "{" ^ fold_left' (fun r s -> r ^ ",\n" ^ s) "" (map ma_of_row d) ^ "}"

let ma_of_dense_sym = ma_of_dense /%/ dense_of_dense_sym;;

match Sys.argv with
| [| "/Users/roland/Logiciel/ocaml/bin/ocaml" |] -> ()
| [| _ |] -> 
    (let source_file = "<stdin>" in
     let source = (*[open_in source_file]*) stdin in
     let dest = stdout (*[open_out (source_file ^ ".sos")]*) in
     let ((p:poly), (vars:string list)) = 
       try Grammar.Entry.parse problem (Stream.of_channel source) 
       with Stdpp.Exc_located (loc,exc) -> (error_msg source_file (loc,exc); exit 1) in
     close_in source;
     prerr_endline ("vars: " ^ fold_left' (fun a b -> a ^ ", " ^ b) "" vars);
     prerr_endline ("poly: " ^ string_of_poly vars p);
     (if (degree p) mod 2 = 1 then
           (prerr_endline "polynomial of odd degree"; exit (-1)));
     let sm = submonos (length vars) p in
     prerr_endline ((length % string_of_int) sm ^ " submonomials: {" ^ (map (string_of_mono vars) % fold_left' (fun a b -> a ^ ", " ^ b) "") sm ^ "}");
     let (k,f0,fs) = sdp_of_sos sm p in
       prerr_endline ("m = " ^ (length % string_of_int) fs);
     match
       (match fs with [] -> Some [] | _ -> solve_sdp (k,f0,fs))
     with
       | Some lambda ->
           let fs = map (dense_sym_of_sparse_sym k) fs in
           let f0 = dense_sym_of_sparse_sym k f0 in
           let d = disturb (one // num_of_int 1000) lambda in
             ignore (find (fun l ->
                             (* prerr_endline ("lambda: " ^ ma_of_row l); *)
                             let q = lin_comb l f0 fs in
                             (* prerr_endline ("q:\n" ^ ma_of_dense_sym q ^ "\n\n");*)
                             match diag q with
                             | Some (l,d) ->
                                 if exists (fun x -> x </ zero) d then (prerr_endline "not psd (negative diag entry)";false) else
                                   ((* prerr_endline ("l:\n" ^ (dense_of_lowtri % ma_of_dense) l);
                                      prerr_endline ("d:\n" ^ (dense_of_diag % ma_of_dense) d); *)
                                    print_endline ("\nSolution:\n" ^ zld_to_string vars sm l d);
                                    true)
                             | None -> prerr_endline "not decomposoable as LDL^T";false
                          ) d);
                    close_out dest
       | None ->
           close_out dest;
           prerr_endline "No solution found."
    )
| _ -> prerr_endline ("usage: " ^ Sys.argv.(0))
